
#include <string.h>
#include "encrypt/aes/aes.h"

#ifndef KEY_ILLEGAL
#define KEY_ILLEGAL (keyLen != 128 && keyLen != 192 && keyLen != 256)
#endif

void AES_ECB_crypt(int mode, uint8_t *input, const uint8_t *key, uint8_t *output)
{
    mbedtls_aes_context ctx;
    uint32_t keyLen = strlen(key);
    keyLen *= 8;

    if(KEY_ILLEGAL) return;

    mbedtls_aes_init(&ctx);

    if(mode == MBEDTLS_AES_ENCRYPT)
        mbedtls_aes_setkey_enc(&ctx, key, keyLen);
    else
        mbedtls_aes_setkey_dec(&ctx, key, keyLen);

    uint8_t * t_in = (uint8_t *)malloc(16 * sizeof(uint8_t));
    uint8_t * t_out = (uint8_t *)malloc(16 * sizeof(uint8_t));
    if(t_in != NULL && t_out != NULL)
    {
        memcpy(t_in, input, 16);
        mbedtls_aes_crypt_ecb(&ctx, mode, t_in, t_out);
        memcpy(output, t_out, 16);
        free(t_in);
        free(t_out);
    }
}

void AES_CBC_crypt(int mode, uint8_t* input,  const uint8_t* iv,const uint8_t* key, uint8_t *output)
{
    mbedtls_aes_context ctx;
    uint32_t keyLen = strlen(key);
    keyLen *= 8;

    if(KEY_ILLEGAL) return;

    if(mode == MBEDTLS_AES_ENCRYPT)
        mbedtls_aes_setkey_enc(&ctx, key, keyLen);
    else
        mbedtls_aes_setkey_dec(&ctx, key, keyLen);

    uint8_t * t_in = (uint8_t *)malloc(16 * sizeof(uint8_t));
    uint8_t * t_out = (uint8_t *)malloc(16 * sizeof(uint8_t));
    uint8_t t_iv[16];
    if(iv == NULL)
    {
        memset(t_iv, 0, 16);
    }else
    {
        int len = strlen(iv);
        if(len > 16)
        {
            memcpy(t_iv, iv, 16);
        }else
        {
            memset(t_iv, 0, 16);
            memcpy(t_iv, iv, len);
        }
    }

    if(t_in != NULL && t_out != NULL)
    {
        memcpy(t_in, input, 16);
        mbedtls_aes_crypt_cbc(&ctx, mode, 16, t_iv, t_in, t_out);
        memcpy(output, t_out, 16);
        free(t_in);
        free(t_out);
    }
}

/**
 * 不定长加密,pkcs5padding
 */
char *AES_PKCS5Padding_Encrypt(int type, const char *in, const uint8_t* iv,const uint8_t *key)
{
    int inLength = (int) strlen(in);//输入的长度

    int remainder = inLength % 16;
    uint8_t *paddingInput;
    int paddingInputLengt = 0;
    if (inLength < 16)
    {
        paddingInput = (uint8_t *) malloc(16);
        memset(paddingInput, 0, sizeof(paddingInput));
        paddingInputLengt = 16;
        int i;
        for (i = 0; i < 16; i++)
        {
            if (i < inLength)
            {
                paddingInput[i] = in[i];
            } else
            {
                paddingInput[i] = HEX[16 - inLength];
            }
        }
    } else
    {
        int group = inLength / 16;
        int size = 16 * (group + 1);
        paddingInput = (uint8_t *) malloc(size);
        paddingInputLengt = size;

        int dif = size - inLength;
        int i;
        for (i = 0; i < size; i++)
        {
            if (i < inLength)
            {
                paddingInput[i] = in[i];
            } else
            {
                if (remainder == 0)
                {
                    //刚好是16倍数,就填充16个16
                    paddingInput[i] = HEX[0];
                } else
                {    //如果不足16位 少多少位就补几个几  如：少4为就补4个4 以此类推
                    paddingInput[i] = HEX[dif];
                }
            }
        }
    }
    int count = paddingInputLengt / 16;
    //开始分段加密
    char *out = (char *) malloc(paddingInputLengt);
    memset(out, 0, sizeof(out));
    int i;
    for (i = 0; i < count; ++i)
    {
        if(type == ECB)
            AES_ECB_crypt(MBEDTLS_AES_ENCRYPT, paddingInput + i * 16, key, out + i * 16);
        else if(type == CBC)
            AES_CBC_crypt(MBEDTLS_AES_ENCRYPT, paddingInput + i * 16, iv, key, out + i * 16);
    }
    char *base64En = b64_encode(out, paddingInputLengt);
    //LOGE(base64En);
    free(paddingInput);
    free(out);
    return base64En;
}

char *AES_PKCS5Padding_Decrypt(int type, const char *in, const uint8_t* iv,const uint8_t *key)
{
    uint8_t *inputDesBase64 = b64_decode(in, strlen(in));
    //base64的长度是原串的4/3
    const size_t inputLength = (strlen(in) / 4) * 3;
    uint8_t *out = malloc(inputLength);
    memset(out, 0, inputLength);
    size_t count = inputLength / 16;
    if (count <= 0)
    {
        count = 1;
    }
    size_t i;
    for (i = 0; i < count; ++i)
    {
        if(type == ECB)
            AES_ECB_crypt(MBEDTLS_AES_DECRYPT, inputDesBase64 + i * 16, key, out + i * 16);
        else if(type == CBC)
            AES_CBC_crypt(MBEDTLS_AES_DECRYPT, inputDesBase64 + i * 16, iv, key, out + i * 16);
    }

    /**
     *  接下来的工作就把末尾的padding去掉
     *  "abcdefghijklmnop\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\0\0\0\0
     *  To "abcdefghijklmnop\n"
     *
     *  "1\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f\x0f"
     *  To "1\n"
     */
    int *result = findPaddingIndex(out, inputLength - 1);
    int offSetIndex = result[0];
    int lastChar = result[1];
    //检查是不是padding的字符,然后去掉
    const size_t noZeroIndex = inputLength - offSetIndex;
    if (lastChar >= 0 && offSetIndex >= 0)
    {
        int success = JNI_TRUE, i;
        for (i = 0; i < lastChar; ++i)
        {
            size_t index = noZeroIndex - lastChar + i;
            if (!HEX[lastChar] == out[index])
            {
                success = JNI_FALSE;
            }
        }
        if (JNI_TRUE == success)
        {
            out[noZeroIndex - lastChar] = '\0';
            memset(out + noZeroIndex - lastChar + 1, 0, lastChar - 1);
        }
    } else
    {
        out[noZeroIndex] = '\0';
    }
    //LOGE("解密结果:");
    //LOGE(out);
    free(inputDesBase64);
    return (char *) out;
}

int *findPaddingIndex(uint8_t *str, size_t length)
{
    int result[] = {-1, -1}, i, k;
    for (i = 0; i < length; ++i)
    {
        char c = str[length - i];
        if ('\0' != c)
        {
            result[0] = i;
            for (k = 0; k < 16; ++k)
            {
                if (HEX[k] == c)
                {
                    if (0 == k)
                    {
                        k = 16;
                    }
                    result[1] = k;
                    return result;
                }
            }
            return result;
        }
    }
    return result;
}
